[SM90] Optimize FP8 MegaMoE decode path with swap-AB, 4WG heuristic and L2 bank swizzle

Summary

Relative to the dev branch, this PR applies three groups of performance optimizations to the SM90 FP8 MegaMoE kernel, covering the L2 epilogue, the 4-WG (split-MN) launch heuristic, and the decode swap-AB path.

Main changes:

• Eliminate shared-memory bank conflicts in the L2 epilogue via an 8-column-granularity XOR swizzle, applied symmetrically on both the STS write and LDS read sides; the permutation is self-cancelling.
• Refine the 4-WG launch threshold to depend on both intermediate_hidden and GPU model (SM count). On H20 + pro (intermediate_hidden >= 3072), where the smaller SM count makes register spilling costly, prefer 2-WG (only launch 4-WG when >512 tokens/expert); on H20 + flash (intermediate_hidden <= 2048), where 2-WG + BLOCK_N=256 beats 4-WG across part of the batch range, prefer 2-WG (4-WG only when >576); on H200/H100 and the H20 mid-range keep the original >=64 threshold and stay 4-WG-first. The BLOCK_N=256 selection for the 2-WG split-N path is also un-gated from intermediate_hidden >= 3072, so flash picks n=256 instead of being forced to n=128.
• Extend the swap-AB decode path from BLOCK_N=128 to BLOCK_N=256 (2-WG N-split, each WG driving two m64n64 slabs), and fix the L1-output TMA descriptor to cover the full tile rather than degrading to wg_l1_out_block_n = 64.

Performance

Notes:

• Baseline: dev branch vs this PR.
• Hardware: 8× H20-3E.
• Two configs: DS-V4 Flash (256 experts) / Pro (384 experts).
• speedup = latency(baseline) / latency(this PR).

Flash shape (256 experts)

Pro shape (384 experts)

Pro sees a ~7–18% latency reduction across t16–t1024 (1.22× peak at t512). Flash also gains across the prefill range: t256 1.09× (BLOCK_N=256 un-gated, was 1.00×), t1024 1.05× and t2048 1.04× (4-WG → 2-WG + n256, t2048 flipping from 0.99× to a gain); t512 dips to 0.98× on the 2-WG/4-WG boundary, where 4-WG still wins — consistent with the 2-WG + n256 advantage being shape-dependent. The decode range keeps ~1.01×/1.08× (t16/t64). t4096+ stays 4-WG and is essentially flat near the FP8 compute ceiling (251 TFLOPS at t8192, ~85% of H20 FP8 peak).

Accuracy

All three changes are performance-only and do not change numerical results; accuracy is consistent with the dev branch:

• L2 epilogue bank swizzle: the 8-column XOR permutation is applied symmetrically on STS/LDS and is self-cancelling, so the column order written to global memory is unchanged.
• 4-WG launch threshold + BLOCK_N=256 selection: only changes which kernel shape / tile granularity is selected per shape — 2-WG n256 (two m64n128 WGMMA halves sharing the A load) vs 4-WG, both tiling the same C block. Each shape independently accumulates the full-tile GEMM, so the result is mathematically equivalent.
• swap-AB extended to BLOCK_N=256: performs the same GEMM with FP32 accumulation and unchanged K-order, splitting the original m64n128 into two m64n64 slabs.

Single-operator and model-level accuracy validations match the dev branch, with no accuracy regression.